#!/beegfs/group/lqcd/software/anaconda3/bin/python3
# encoding: utf-8

import os
import sys
from math import *
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt

#user parameters
Nt = 9
T = 5
dt= T/(Nt+1)
xmax = 5
Nx = 16



device = sys.argv[1]
mode = sys.argv[2]

sbatch_base="#!/bin/bash\n\
#SBATCH -J hrmnc_oscltr\n\
#SBATCH -p debug\n\
#SBATCH --nodes=1\n\
#SBATCH --ntasks-per-node=1\n\
#SBATCH --gres=gpu:1\n\
#SBATCH --get-user-env\n\
#SBATCH --time=00:10:00\n"


if device =="CPU":
    Npnts = 200000
elif device =="GPU":
    Npnts = 2048 # gVegas: Npnts * 1K in gVegas

if mode == "compile":
#    os.system("chmod +x ./"+device+"/compile.sh")
    os.system("cd ./"+device+" && ./compile.sh")

if mode == "run":
    if device == "CPU":
        sbmt_scrpt = sbatch_base + "\n#SBATCH --output="+device+".out\n#SBATCH --error="+device+".err\n" \
                    +"srun --nodes=1 --ntasks-per-node=1 --time=00:10:00 ./"+device+"/hrmnc_oscltr.o " \
                    +str(Nt)+" "+str(dt)+" "+str(xmax)+" "+str(Nx)+" "+str(Npnts)
    elif device == "GPU":
        sbmt_scrpt = sbatch_base + "\n#SBATCH --output="+device+".out\n#SBATCH --error="+device+".err\n" \
                    "srun --nodes=1 --ntasks-per-node=1 --time=00:10:00 ./"+device+"/hrmnc_oscltr.o " \
                    +str(Nt)+" "+str(dt)+" "+str(xmax)+" "+str(Nx)+" "+str(Npnts)
    

    sh = open("submit_"+device+".sh", "w", encoding="utf-8")
    sh.write(sbmt_scrpt)
    sh.close()

    os.system("sbatch ./submit_"+device+".sh")    
   

if mode == "analysis":
    df = pd.read_csv("./"+device+"_hrmnc_oscltr.dat",names =["x","res","err","chisq"]) 
    x_lst = df["x"].values
    cntrl = df["res"].values
    err = df["err"].values

    fig = plt.figure()
    plt.errorbar(x_lst,cntrl,yerr=err,markersize=4.,fmt="o",ecolor="r",color="b",elinewidth=2)
    plt.yscale("log")
    plt.grid(True)

    x_lnspc = np.linspace(min(x_lst),max(x_lst),61)
    psi_0 = lambda x: pow(pi,-1/4)*exp(-x*x/2)
    y_lst = [psi_0(x)**2*exp(-T/2) for x in x_lnspc]
    plt.plot(x_lnspc,y_lst,color="r")
    plt.yscale("log")
    plt.grid(True)

    plt.savefig(device+"_hrmnc_oscltr.png")

if mode == "clean":
    os.system("cd ./"+device+" && rm *.o")
    os.system("rm *.err *.out")
    os.system("rm ./"+device+"_hrmnc_oscltr.dat")
    os.system("rm ./"+device+"_hrmnc_oscltr.png")
